import os
import json
import uuid
import requests
from typing import List, Dict
from dataclasses import dataclass
import hashlib
import openai  # 我们将用它调用兼容 OpenAI 的服务
from neo4j import GraphDatabase
'''
此模块提供了一个函数，用于从文本中提取问题和回答，
并将这些信息导入到Neo4j数据库中。
'''
# =======================
# 1. 配置区（请根据你的环境修改）
# =======================

# 🔧 OpenAI 兼容 API 配置（例如：阿里云百炼、vLLM 等）
OPENAI_BASE_URL = "https://api.siliconflow.cn/v1"  # 替换为你的实际地址
OPENAI_API_KEY = "sk-xpczcmiojslifyfbfzkrkctbgbilljxoekpevjhcdnrcdzjb"                       # 替换为你的密钥
LLM_MODEL_NAME = "Qwen/Qwen3-235B-A22B-Instruct-2507"                           # 模型名（由平台提供）
EMBEDDING_MODEL_NAME = "BAAI/bge-m3"            # 嵌入模型名

# 🛢️ Neo4j 配置
NEO4J_URI = "bolt://localhost:7687"
NEO4J_USER = "neo4j"
NEO4J_PASSWORD = "12345678"

# 初始化 OpenAI 客户端（兼容 OpenAI 格式的任意平台）
client = openai.OpenAI(base_url=OPENAI_BASE_URL, api_key=OPENAI_API_KEY)

# 初始化 Neo4j 连接
driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
driver.verify_connectivity()


# =======================
# 2. 数据结构定义
# =======================
@dataclass
class QA:
    question: str
    answer: str


@dataclass
class Entity:
    name: str
    type: str


@dataclass
class Relation:
    from_entity: str
    to_entity: str
    relation: str


# =======================
# 3. 步骤1：用 LLM 分离问题和回答
# =======================
def extract_question_answer(raw_text: str) -> QA:
    prompt = f"""
    你是一个运维知识结构化引擎，请从以下文本中提取出“问题”和“回答”。
    必须完全符合原文
    请以 JSON 格式输出：
    {{
      "question": "问题文本",
      "answer": "回答文本"
    }}

    输入文本：
    {raw_text}
    """

    response = client.chat.completions.create(
        model=LLM_MODEL_NAME,
        messages=[{"role": "user", "content": prompt}],
        response_format={"type": "json_object"},
        temperature=0,
        seed=1
    )

    content = response.choices[0].message.content
    try:
        data = json.loads(content)
        return QA(question=data["question"].strip(), answer=data["answer"].strip())
    except Exception as e:
        raise ValueError(f"无法解析 LLM 输出: {content}") from e


# =======================
# 4. 步骤2：用 LLM 抽取实体和关系
# =======================
def extract_entities_and_relations(answer_text: str) -> tuple[List[Entity], List[Relation]]:
    prompt = f"""
    请从以下运维回答中提取：
    1. 所有技术实体（如服务、服务器、IP、端口、错误码、配置文件等），并标注类型。
    2. 实体之间的语义关系（如“依赖”、“监听”、“导致”、“配置于”等动词关系）。

    输出为 JSON：
    {{
      "entities": [
        {{"name": "nginx", "type": "Service"}},
        {{"name": "80", "type": "Port"}}
      ],
      "relations": [
        {{"from": "nginx", "to": "80", "relation": "listens_on"}}
      ]
    }}

    回答内容：
    {answer_text}
    """

    response = client.chat.completions.create(
        model=LLM_MODEL_NAME,
        messages=[{"role": "user", "content": prompt}],
        response_format={"type": "json_object"},
        temperature=0,
        seed=1
    )

    content = response.choices[0].message.content
    try:
        data = json.loads(content)
        entities = [Entity(name=e["name"], type=e["type"]) for e in data.get("entities", [])]
        relations = [
            Relation(from_entity=r["from"], to_entity=r["to"], relation=r["relation"])
            for r in data.get("relations", [])
        ]
        return entities, relations
    except Exception as e:
        raise ValueError(f"无法解析实体关系输出: {content}") from e


# =======================
# 5. 步骤3：调用嵌入模型生成向量
# =======================
def get_embedding(text: str) -> List[float]:
    response = client.embeddings.create(
        model=EMBEDDING_MODEL_NAME,
        input=text
    )
    return response.data[0].embedding


# =======================
# 6. 步骤4：写入 Neo4j 图数据库
# =======================
def save_to_neo4j(qa: QA, entities: List[Entity], relations: List[Relation]):
    # 用问题+回答的内容生成确定性 ID
    content_key = f"{qa.question} || {qa.answer}"
    doc_id = generate_deterministic_id(content_key)
    ans_id = generate_deterministic_id(content_key)  # 同一个 QA 对，Answer 和 Document 共享 ID 或也可分开

    # 生成嵌入向量（相同文本 → 相同 embedding，前提是模型确定）
    doc_emb = get_embedding(qa.question)
    ans_emb = get_embedding(qa.answer)

    # 准备数据
    entities_data = [{"name": e.name, "type": e.type} for e in entities]
    relations_data = [
        {"from": r.from_entity, "to": r.to_entity, "relation": r.relation}
        for r in relations
    ]

    with driver.session() as session:
        # 使用 MERGE 而不是 CREATE，确保不会重复插入
        result = session.run("""
        // 步骤1：MERGE Document 和 Answer 节点（基于确定性 ID）
        MERGE (d:Document {
          id: $doc_id
        })
        ON CREATE SET d.text = $question, d.type = "question", d.embedding = $doc_emb
        ON MATCH SET d.text = $question, d.embedding = $doc_emb  // 若已存在则更新内容（可选）

        MERGE (a:Answer {
          id: $ans_id
        })
        ON CREATE SET a.text = $answer, a.embedding = $ans_emb
        ON MATCH SET a.text = $answer, a.embedding = $ans_emb

        // 建立 HAS_ANSWER 关系（使用 MERGE 防止重复）
        MERGE (d)-[r_has:HAS_ANSWER]->(a)

        // 传参继续
        WITH d, a, $entities AS entities, $relations AS relations

        // 步骤2：MERGE 所有实体（基于 name 唯一性）
        UNWIND entities AS entity
        MERGE (e:Entity {name: entity.name})
        SET e.type = COALESCE(e.type, entity.type)  // 保留已有 type，若无则设置

        // 创建 Answer-MENTIONS_ENTITY->Entity 关系（用 MERGE 避免重复）
        MERGE (a)-[r_mention:MENTIONS_ENTITY]->(e)

        WITH d, a, COLLECT(e) AS e_list, relations

        // 步骤3：处理实体间的关系（from -> to + relation 类型）
        UNWIND relations AS rel
        MATCH (from:Entity {name: rel.from})
        MATCH (to:Entity {name: rel.to})
        // 注意：关系类型是动态的，不能直接用变量做关系类型，需用 apoc 或字符串执行
        // 我们这里改用通用关系 + 属性方式存储 relation 类型
        MERGE (from)-[rel_edge:RELATED_TO]->(to)
        SET rel_edge.relation = rel.relation  // 去重后保留最后的 relation 类型（或不做 set）

        RETURN d.id AS document_id, a.id AS answer_id
        """,
        doc_id=doc_id,
        ans_id=ans_id,
        question=qa.question,
        answer=qa.answer,
        doc_emb=doc_emb,
        ans_emb=ans_emb,
        entities=entities_data,
        relations=relations_data
        )

        record = result.single()
        print(f"✅ 成功写入/合并：Document ID = {record['document_id']}")

# =======================
# 7. 主流程（从文件读取 + 按空行分割）
# =======================
def read_qa_blocks_from_file(file_path: str) -> List[str]:
    """
    读取文本文件，按空行分割成多个文本块（每个块是一组 QA 内容）
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"文件未找到: {file_path}")

    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read().strip()

    # 使用空行分割（支持 \n\n 或多个换行）
    blocks = [block.strip() for block in content.split('\n\n') if block.strip()]
    return blocks
# =======================
# 新增：根据字符串内容生成确定性 ID
# =======================
def generate_deterministic_id(content: str) -> str:
    """基于 SHA256 哈希生成唯一、可复现的 ID"""
    return hashlib.sha256(content.encode('utf-8')).hexdigest()

# =======================
# 8. 封装主流程为函数
# =======================
def process_qa_file(file_path: str = "qa_input.txt"):
    print(f"📄 正在读取文件: {file_path}")
    try:
        qa_blocks = read_qa_blocks_from_file(file_path)
        print(f"✅ 共读取到 {len(qa_blocks)} 个问答块\n")
    except Exception as e:
        print(f"❌ 读取文件失败: {e}")
        exit(1)

    # 处理每一个块
    for i, raw_text in enumerate(qa_blocks, start=1):
        print(f"📌 处理第 {i}/{len(qa_blocks)} 个问答块...")
        try:
            print("🔍 正在提取问题和回答...")
            qa = extract_question_answer(raw_text)
            print("📝 问题:", qa.question)
            print("💡 回答:", qa.answer)

            print("\n🔍 正在抽取实体和关系...")
            entities, relations = extract_entities_and_relations(qa.answer)
            print("🏷️  实体:", [(e.name, e.type) for e in entities])
            print("🔗 关系:", [(r.from_entity, r.relation, r.to_entity) for r in relations])

            print("\n🧠 正在生成 embedding 并写入 Neo4j...")
            save_to_neo4j(qa, entities, relations)

            print(f"✅ 第 {i} 个块处理完成\n" + "="*50 + "\n")

        except Exception as e:
            print(f"❌ 第 {i} 个块处理失败: {e}\n")

    print("🎉 所有内容处理完成！访问 http://localhost:7474 查看图谱数据")

if __name__ == "__main__":
    process_qa_file()